#include "imageannotation.h"

using namespace std;
using namespace cv;

bool ImageAnnotation::read()
{
    readFromDat();

    return true;
}

bool ImageAnnotation::write()
{
    write2dat();
    write2xml();

    return true;
}

void ImageAnnotation::Add(const ImageAnnEntry &annEntry)
{
    imageAnnEntries.push_back(annEntry);
    annEntryCounts = imageAnnEntries.size();
}

void ImageAnnotation::Remove(const ImageAnnEntry &annEntry)
{
    imageAnnEntries.erase(remove(imageAnnEntries.begin(), imageAnnEntries.end(), annEntry), imageAnnEntries.end());
    // 更新数量
    annEntryCounts = imageAnnEntries.size();
}

void ImageAnnotation::clear()
{
    imageAnnEntries.clear();
    annEntryCounts = 0;
}

bool ImageAnnotation::write2dat()
{
    //打开文件
    ofstream out(annFileName, ios::trunc);
    if (!out)
    {
        return false;
    }
    // 写入数据文件对应的图片文件名
    out << imageFileName << endl;
    // 写入标注数量
    out << annEntryCounts << endl;
    // 将图片标注条目写入到文件中

    for (auto entry : imageAnnEntries)
    {
        // 写入条目名称
        out << endl;
        out << entry.getEntryName() << endl;
        // 写入当前标注条目下标注项个数
        out << entry.getItemCounts() << endl;

        for(auto item :entry.getImageItems())
        {
            // 写入类型
            Type t = item.getAnnType();
            if(t == Empty_Type)
                continue;
            out << endl;
            out << t << endl;
            // 写入点的数据
            for(auto p : item.getPoints())
            {
                out << p.x << "\t" << p.y << endl;
            }
        }
    }

    out << endl;
    out.close();

    return true;
}

bool ImageAnnotation::write2xml()
{
    string xmlFileName = imageFileName + ".xml";
    size_t pos = imageFileName.find_last_of('/');
    string imageName(imageFileName.begin() + pos + 1, imageFileName.end());
    string newPath(imageFileName.begin(), imageFileName.begin() + pos);
    size_t pos2 = newPath.find_last_of('/');
    string floder(newPath.begin() + pos2 + 1, newPath.end());

    Mat src = imread(imageFileName);
    QFile file(string2QString(xmlFileName));
    file.open(QIODevice::WriteOnly);
    QXmlStreamWriter xmlWriter(&file);
    xmlWriter.setAutoFormatting(true);
    xmlWriter.writeStartDocument();
    xmlWriter.writeStartElement("annotation");
    xmlWriter.writeTextElement("floder",string2QString(floder));
    xmlWriter.writeTextElement("filename", string2QString(imageName));
    xmlWriter.writeStartElement("source");
    xmlWriter.writeStartElement("database");
    xmlWriter.writeEndElement();
    xmlWriter.writeStartElement("annotation");
    xmlWriter.writeEndElement();
    xmlWriter.writeStartElement("image");
    xmlWriter.writeEndElement();
    xmlWriter.writeStartElement("flickrid");
    xmlWriter.writeEndElement();
    xmlWriter.writeEndElement();
    xmlWriter.writeStartElement("size");
    xmlWriter.writeTextElement("width", QString::number(src.cols));
    xmlWriter.writeTextElement("height", QString::number(src.rows));
    xmlWriter.writeTextElement("depth", QString::number(src.channels()));
    xmlWriter.writeEndElement();
    xmlWriter.writeTextElement("segmented", QString::number(0));
    for(auto entry : imageAnnEntries)
    {
        for(auto item : entry.getImageItems())
        {
            if(item.getAnnType() != Rectangle_Type)
                continue;
            xmlWriter.writeStartElement("object");
            xmlWriter.writeTextElement("name", QString::fromStdString(entry.getEntryName()));
            xmlWriter.writeTextElement("pose", QString("Unspecified"));
            xmlWriter.writeTextElement("truncated", QString::number(0));
            xmlWriter.writeTextElement("difficult", QString::number(0));
            xmlWriter.writeStartElement("bndbox");
            vector<Point> pts = item.getPoints();
            xmlWriter.writeTextElement("xmin", QString::number(pts[0].x));
            xmlWriter.writeTextElement("ymin", QString::number(pts[0].y));
            xmlWriter.writeTextElement("xmax", QString::number(pts[2].x));
            xmlWriter.writeTextElement("ymax", QString::number(pts[2].y));
            xmlWriter.writeEndElement();
            xmlWriter.writeEndElement();
        }
    }
    xmlWriter.writeEndElement();

    file.close();
    return true;
}

bool ImageAnnotation::readFromDat()
{
    ifstream in(annFileName);
    if (!in)
        return false;

    string line;
    stringstream ss;
    // 图片文件名(丢掉，没用)
    getline(in, line);
    //imageFileName = line;
    // 图片标注条目数量
    getline(in, line);
    annEntryCounts = stoi(line.c_str());
    getline(in, line); // 丢掉一个空行

    ImageAnnEntry entry;
    ImageAnnItem item;
    cv::Point p;

    for(int i = 0; i < annEntryCounts; ++i)
    {
        // 标注条目名称
        getline(in, line);
        entry.clear();
        entry.setEntryName(line);
        entry.setColor(Scalar(255, 0, 255));
        // 标注条目中标注项个数
        getline(in, line);
        int itemCount = stoi(line.c_str());
        entry.setItemCounts(itemCount);
        // 丢掉空行
        getline(in, line);
        // 开始读取标注项
        for(int j = 0; j < itemCount; ++j)
        {
            // 读取标注项类型
            getline(in, line);
            Type t = static_cast<Type>(stoi(line.c_str()));
            item.clear();
            item.setAnnType(t);
            // 读取改类型下的点
            for(getline(in, line); line != ""; getline(in, line))
            {
                ss << line;
                ss >> p.x >> p.y;
                item.addPoint(p);
                ss.clear();
            }
            entry.addItem(item);
        }
        imageAnnEntries.push_back(entry);
    }

    in.close();

    return true;
}
